package com.zyk.scaffold.gateway.filter;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONValidator;
import com.zyk.scaffold.core.component.SimpleCountRedisLimitComponent;
import com.zyk.scaffold.core.context.AppContext;
import com.zyk.scaffold.core.context.JwtHelper;
import com.zyk.scaffold.common.exception.BusinessException;
import com.zyk.scaffold.common.utils.Md5SignUtils;
import com.zyk.scaffold.gateway.constant.GatewayConstants;
import com.zyk.scaffold.gateway.enums.GatewayErrorCodeEnum;
import com.zyk.scaffold.gateway.properties.OpenApiProperties;
import com.zyk.scaffold.gateway.support.HttpMessageReaderMultipartFormDataSupport;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.ResolvableType;
import org.springframework.core.codec.Hints;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.ReactiveHttpInputMessage;
import org.springframework.http.codec.multipart.FilePart;
import org.springframework.http.codec.multipart.FormFieldPart;
import org.springframework.http.codec.multipart.Part;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.*;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Supplier;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

@Component
public class OpenApiFilter implements GlobalFilter, Ordered {
    private static final Logger log = LoggerFactory.getLogger(OpenApiFilter.class);

    @Autowired
    private OpenApiProperties openApiProperties;

    @Autowired
    private SimpleCountRedisLimitComponent redisLimitComponent;

    private final HttpMessageReaderMultipartFormDataSupport partReader = new HttpMessageReaderMultipartFormDataSupport();

    private static final ResolvableType FORM_DATA_TYPE = ResolvableType.forClassWithGenerics(MultiValueMap.class, String.class, Part.class);

    private static final Mono<? extends LinkedMultiValueMap<String, Part>> EMPTY_MULTIPART_DATA = Mono
            .just(new LinkedMultiValueMap<String, Part>(0)).cache();

    private static final Mono<HashMap<String, Object>> EMPTY_HASH_MAP = Mono.just(new HashMap<String, Object>(0)).cache();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        Optional<String> appIdOptional = getAppId(request);
        String appId = appIdOptional.orElse(null);
        if (null == appId) {
            return chain.filter(exchange);
        }

        String path = request.getURI().getRawPath();
        OpenApiProperties.App app = openApiProperties.getAppByAppId(appId).orElse(null);
        if (null == app) {
            throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_APP_NOT_AVAILABLE);
        }
        if (isOpenApi(app.getUris(), path)) {
            Long timeStamp = getTimeStamp(request).orElse(null);
            if (null == timeStamp) {
                throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_TIME_STAMP_REQUIRED);
            }
            List<OpenApiProperties.Limit> limits = app.getLimits();
            if (CollectionUtils.isNotEmpty(limits)) {
                verifLimit(limits, appId, path);
            }
            if (app.getEnableTimeOut() && isSignTimeOut(timeStamp, app.getTimeoffset(), app.getTimeout())) {
                log.info("timeStamp:{}, app:{}, time:{}", timeStamp, JSON.toJSONString(app), LocalDateTime.now().toInstant(ZoneOffset.ofHours(8)).toEpochMilli());
                throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_SIGN_TIME_OUT);
            }
            String sign = getSign(request).orElse(null);
            if (null == sign) {
                throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_SIGN_REQUIRED);
            }
            String appKey = app.getAppKey();
            if (null == appKey) {
                throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_APP_KEY_NOT_AVAILABLE);
            }
            if ((request.getMethod() == HttpMethod.POST
                    || request.getMethod() == HttpMethod.PUT
                    || request.getMethod() == HttpMethod.PATCH)) {
                return operationPostExchange(exchange, chain, app, sign, timeStamp);
            }
            if ((request.getMethod() == HttpMethod.GET)) {
                return operationGetExchange(exchange, chain, app, sign, timeStamp);
            }
        }
        return chain.filter(exchange);
    }

    private Mono<Void> operationPostExchange(ServerWebExchange exchange, GatewayFilterChain chain, OpenApiProperties.App app, String sign, Long timeStamp) {
        String appId = app.getAppId();
        String appKey = app.getAppKey();
        ServerHttpRequest request = exchange.getRequest();
        return DataBufferUtils.join(request.getBody())
                .flatMap(dataBuffer -> Mono.just(Optional.of(dataBuffer)))
                .defaultIfEmpty(Optional.empty())
                .flatMap(dataBufferOptional -> {
                    DataBuffer dataBuffer = dataBufferOptional.orElse(null);
                    int byteCount = dataBuffer == null ? 0 : dataBuffer.readableByteCount();
                    byte[] bytes = new byte[byteCount];
                    if (dataBuffer != null) {
                        dataBuffer.read(bytes);
                        DataBufferUtils.release(dataBuffer);
                    }
                    MediaType contentType = request.getHeaders().getContentType();
                    long contentLength = request.getHeaders().getContentLength();
                    Mono<Map<String, Object>> mapParams;
                    Supplier<Mono<DataBuffer>> supplierDataBuffer = () -> Mono.just(DataBufferUtils.retain(exchange.getResponse().bufferFactory().wrap(bytes)));
                    Supplier<ServerHttpRequest> supplierNewReq = () -> new ServerHttpRequestDecorator(request) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return Flux.defer(supplierDataBuffer);
                        }
                    };
                    BiFunction<MediaType, String, Boolean> startsWithContentType = ((contentTypeValue, mediaType) -> contentTypeValue != null && contentTypeValue.toString().startsWith(mediaType));
                    if (contentLength == 0 || bytes.length == 0) {
                        mapParams = Mono.just(new HashMap<>(0));
                    } else if (MediaType.APPLICATION_JSON.equals(contentType) || startsWithContentType.apply(contentType, MediaType.APPLICATION_JSON_VALUE)) {
                        mapParams = Mono.fromCallable(() -> {
                                    String s = new String(bytes, StandardCharsets.UTF_8);
                                    return isJsonObject(s) ? JSONObject.parseObject(s).getInnerMap() : new HashMap<>();
                                }
                        );
                    } else if (MediaType.APPLICATION_FORM_URLENCODED.equals(contentType) || startsWithContentType.apply(contentType, MediaType.APPLICATION_FORM_URLENCODED_VALUE)) {
                        mapParams = Mono.fromCallable(() -> readFormQueryData(parseFormData(new String(bytes, StandardCharsets.UTF_8))));
                    } else if (MediaType.MULTIPART_FORM_DATA.equals(contentType) || startsWithContentType.apply(contentType, MediaType.MULTIPART_FORM_DATA_VALUE)) {
                        mapParams = readMultiFormDataMono(exchange, supplierNewReq.get());
                    } else {
                        return Mono.error(new BusinessException(GatewayErrorCodeEnum.OPEN_API_CONTENT_TYPE_UNSUPPORT));
                    }
                    return mapParams.switchIfEmpty(EMPTY_HASH_MAP)
                            .doOnNext(verifSignConsumer(contentType, appKey, sign, timeStamp, exchange))
                            .then(Mono.defer(() -> chain.filter(
                                    buildOpenApiHeader(app, exchange, creatAppToken(app, sign, timeStamp))
                                            .mutate()
                                            .request(supplierNewReq.get())
                                            .build()
                            )));
                });
    }

    private static boolean isJsonObject(String rawConfig) {
        try {
            JSONValidator validator = JSONValidator.from(rawConfig);
            return validator.validate() && validator.getType() == JSONValidator.Type.Object;
        } catch (Exception e) {
            // ignore exception and return false
        }
        return false;
    }

    private Mono<Void> operationGetExchange(ServerWebExchange exchange, GatewayFilterChain chain, OpenApiProperties.App app, String sign, Long timeStamp) {
        String appKey = app.getAppKey();
        MultiValueMap<String, String> queryParams = exchange.getRequest().getQueryParams();
        Map<String, Object> mapParams = readFormQueryData(queryParams);
        log.info("get request sign data:{}", mapParams);

        if (!Md5SignUtils.verifSignMap(mapParams, appKey, sign, timeStamp)) {
            throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_SIGN_ERROR);
        }
        exchange.getAttributes().put(GatewayConstants.OPEN_API_KEY_AUTH, Boolean.TRUE);
        exchange = buildOpenApiHeader(app, exchange, creatAppToken(app, sign, timeStamp));
        return chain.filter(exchange);
    }

    private Mono<Map<String, Object>> readMultiFormDataMono(ServerWebExchange exchange, ReactiveHttpInputMessage message) {
        return this.partReader
                .read(FORM_DATA_TYPE, message, Hints.from(Hints.LOG_PREFIX_HINT, exchange.getLogPrefix()))
                .collectMultimap(Part::name)
                .map(this::toMultiValueMap)
                .switchIfEmpty(EMPTY_MULTIPART_DATA)
                .map(this::readMultiFormData)
                .cache();
    }

    private Consumer<? super Map<String, Object>> verifSignConsumer(MediaType contentType, String appKey, String sign, Long timeStamp, ServerWebExchange exchange) {
        return map -> {
            log.info("contentType:{}, sign data:{}", contentType, map);
            if (!Md5SignUtils.verifSignMap(map, appKey, sign, timeStamp)) {
                throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_SIGN_ERROR);
            }
            exchange.getAttributes().put(GatewayConstants.OPEN_API_KEY_AUTH, Boolean.TRUE);
        };
    }

    private boolean verifLimit(List<OpenApiProperties.Limit> limits, String appId, String path) {
        AntPathMatcher antPathMatcher = new AntPathMatcher();
        HashSet<String> matchingSet = new HashSet<>();
        for (OpenApiProperties.Limit limit : limits) {
            String matching = limit.getMatching();
            if (matchingSet.add(matching) && antPathMatcher.match(matching, path)) {
                int frequency = limit.getFrequency();
                if (frequency < 0) {
                    continue;
                }
                if (frequency > 0 && redisLimitComponent.tryAccess("open-api:limit:" + appId + ":" + matching, frequency, limit.getTimeWindow())) {
                    continue;
                }
                throw new BusinessException(GatewayErrorCodeEnum.OPEN_API_FREQUENCY_OUT_OF_LIMIT);
            }
        }
        return Boolean.TRUE;
    }


    private String creatAppToken(OpenApiProperties.App app, String sign, Long timeStamp) {
        AppContext appContext = new AppContext();
        appContext.setAppId(app.getAppId());
        appContext.setAppName(app.getAppName());
        appContext.setSign(sign);
        appContext.setTimeStamp(timeStamp);
        appContext.setTenantCode(app.getTenantCode());
        appContext.setTenantName(app.getTenantName());
        appContext.setExpire(120);
        String jwt = JwtHelper.toJwt(appContext);
        return jwt;
    }

    private ServerWebExchange buildOpenApiHeader(OpenApiProperties.App app, ServerWebExchange exchange, String appToken) {
        return exchange.mutate()
                .request(exchange
                        .getRequest()
                        .mutate()
                        .headers(httpHeaders -> httpHeaders.remove(GatewayConstants.OPEN_API_TOKEN_HEADER))
                        .header(GatewayConstants.OPEN_API_TOKEN_HEADER, appToken)
                        .header(GatewayConstants.GATEWAY_FILTER_HEADER_TENANT_CODE, app.getTenantCode())
                        .build()
                ).build();
    }

    private Map<String, Object> readFormQueryData(MultiValueMap<String, String> stringStringMultiValueMap) {
        HashMap<String, Object> stringObjectHashMap = new HashMap<>();
        Set<String> strings = stringStringMultiValueMap.keySet();
        for (String key : strings) {
            stringObjectHashMap.put(key, stringStringMultiValueMap.getFirst(key));
        }
        return stringObjectHashMap;
    }

    private Map<String, Object> readMultiFormData(MultiValueMap<String, Part> formData) {
        HashMap<String, Object> stringObjectHashMap = new HashMap<>();
        Set<String> strings = formData.keySet();
        for (String key : strings) {
            Part part = formData.getFirst(key);
            if (part instanceof FormFieldPart) {
                stringObjectHashMap.put(key, ((FormFieldPart) part).value());
            }
            if (part instanceof FilePart) {
                log.info("Multi Form File Part Name:{}", ((FilePart) part).filename());
            }
        }
        return stringObjectHashMap;
    }

    private MultiValueMap<String, String> parseFormData(String body) {
        String[] pairs = org.springframework.util.StringUtils.tokenizeToStringArray(body, "&");
        MultiValueMap<String, String> result = new LinkedMultiValueMap<>(pairs.length);
        try {
            for (String pair : pairs) {
                int idx = pair.indexOf('=');
                if (idx == -1) {
                    result.add(URLDecoder.decode(pair, StandardCharsets.UTF_8.name()), null);
                } else {
                    String name = URLDecoder.decode(pair.substring(0, idx), StandardCharsets.UTF_8.name());
                    String value = URLDecoder.decode(pair.substring(idx + 1), StandardCharsets.UTF_8.name());
                    result.add(name, value);
                }
            }
        } catch (UnsupportedEncodingException ex) {
            throw new IllegalStateException(ex);
        }
        return result;
    }

    private boolean isSignTimeOut(Long timeStamp, Long timeoffset, Long timeout) {
        long curTimeStampOffset = LocalDateTime.now().toInstant(ZoneOffset.ofHours(8)).toEpochMilli() + timeoffset;
        return !(curTimeStampOffset >= timeStamp && curTimeStampOffset <= timeStamp + timeout);
    }

    private boolean isOpenApi(List<String> uris, String path) {
        if(CollectionUtils.isEmpty(uris)){
            return Boolean.FALSE;
        }
        for (String uri : uris) {
            if (Pattern.matches(uri, path)) {
                return Boolean.TRUE;
            }
        }
        return Boolean.FALSE;
    }

    private Optional<String> getAppId(ServerHttpRequest request) {
        List<String> appIds = request.getHeaders().get(GatewayConstants.OpenApiHeader.APP_ID.getKey());
        if (CollectionUtils.isEmpty(appIds)) {
            return Optional.empty();
        }
        return Optional.ofNullable(appIds.get(0));
    }

    private Optional<String> getSign(ServerHttpRequest request) {
        List<String> appSigns = request.getHeaders().get(GatewayConstants.OpenApiHeader.API_SIGN.getKey());
        if (CollectionUtils.isEmpty(appSigns)) {
            return Optional.empty();
        }
        return Optional.ofNullable(appSigns.get(0));
    }

    private Optional<Long> getTimeStamp(ServerHttpRequest request) {
        List<String> timeStamps = request.getHeaders().get(GatewayConstants.OpenApiHeader.API_TIME_STAMP.getKey());
        if (CollectionUtils.isEmpty(timeStamps)) {
            return Optional.empty();
        }
        long timeStamp = NumberUtils.toLong(timeStamps.get(0));
        if (timeStamp == 0) {
            return Optional.empty();
        }
        return Optional.of(timeStamp);
    }

    private LinkedMultiValueMap<String, Part> toMultiValueMap(Map<String, Collection<Part>> map) {
        return new LinkedMultiValueMap<>(map.entrySet().stream()
                .collect(Collectors.toMap(Map.Entry::getKey, e -> toList(e.getValue()))));
    }

    private List<Part> toList(Collection<Part> collection) {
        return collection instanceof List ? (List<Part>) collection : new ArrayList<>(collection);
    }

    @Override
    public int getOrder() {
        return Integer.MIN_VALUE + 1000;
    }
}
